import asyncio
import ctypes

import cv2
import mss
import numpy as np
import websockets
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack, RTCIceCandidate, MediaStreamTrack
from av import VideoFrame
import fractions
import json




class CustomVideoStreamTrack(VideoStreamTrack): # 获取摄像头流代码
    def __init__(self, camera_id):
        super().__init__()
        self.cap = cv2.VideoCapture(camera_id)
        self.frame_count = 0

    async def recv(self):
        self.frame_count += 1
        print(f"Sending frame {self.frame_count}")
        ret, frame = self.cap.read()
        if not ret:
            print("Failed to read frame from camera")
            return None
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video_frame = VideoFrame.from_ndarray(frame, format="rgb24")
        video_frame.pts = self.frame_count
        video_frame.time_base = fractions.Fraction(1, 30)  # Use fractions for time_base
        return video_frame


async def on_ice_candidate(pc, candidate):
    print(f"ICE candidate: {candidate}")

async def setup_webrtc_and_run(ws_url, camera_id):
    pc = RTCPeerConnection()
    video_sender = CustomVideoStreamTrack(camera_id)

    pc.addTrack(video_sender)

    try:
        async with websockets.connect(ws_url) as ws:
            @pc.on("datachannel")
            def on_datachannel(channel):
                print(f"Data channel established: {channel.label}")

            @pc.on("connectionstatechange")
            async def on_connectionstatechange():
                print(f"Connection state is {pc.connectionState}")
                if pc.connectionState == "connected":
                    print("WebRTC connection established successfully")

            @pc.on("icecandidate")
            async def on_icecandidate(candidate):
                if candidate:
                    await on_ice_candidate(pc, candidate)

            offer = await pc.createOffer()
            await pc.setLocalDescription(offer)
            print(f"Sending offer: {offer}")
            await ws.send(json.dumps({'sdp': pc.localDescription.sdp, 'type': pc.localDescription.type}))

            # 等待接收answer
            answer = await ws.recv()
            answer_dict = json.loads(answer)
            print(f"Received answer: {answer_dict}")
            remote_desc = RTCSessionDescription(sdp=answer_dict['sdp'], type=answer_dict['type'])
            await pc.setRemoteDescription(remote_desc)

    except Exception as e:
        print(f"An error occurred with WebSocket: {e}")
    finally:
        # 注意：不要在这里关闭WebRTC连接
        print("WebSocket connection closed, but WebRTC connection remains open")

async def main():
    ws_url = "ws://localhost:8765"  # WebSocket URL of the signaling server
    camera_id = 0  # Change this to the appropriate camera ID
    await setup_webrtc_and_run(ws_url, camera_id)
    while True:
        await asyncio.sleep(1)
    # 等待WebRTC连接关闭（例如，用户请求断开连接）
    # 在这里你可以添加逻辑来等待用户输入或其他事件来关闭WebRTC连接
    # 例如:
    # await asyncio.sleep(1000)  # 等待一段时间或用户输入
    # await pc.close()

if __name__ == "__main__":
    asyncio.run(main())
